import cv2
import torch

# 加载模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='training/yolov5s.pt')

# 打开摄像头
cap = cv2.VideoCapture(0)  # 0 代表默认摄像头

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 模型推理
    results = model(frame)

    # 将结果渲染到图像上
    annotated_frame = results.render()[0]

    # 显示窗口
    cv2.imshow("YOLOv5 Detection", annotated_frame)

    # 按 'q' 退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()